import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

df = pd.read_csv('eden_experiment_298.csv')
df.fillna(method="bfill", inplace=True)

attention_df = df[df["attention_choice"] == "D"]

abstract_df = attention_df[df["condition (Expert = Description, Abstract = Experience)"] == "Abstract"]
expert_df = attention_df[df["condition (Expert = Description, Abstract = Experience)"] == "Expert"]

abstract_df_by_round = abstract_df.groupby(["Round"]).agg({"Choice": "mean"})
abstract_df_by_round = abstract_df_by_round.groupby([np.arange(200)//5]).agg({"Choice": "mean"})
abstract_df_by_round["Condition"] = "Experience"
abstract_df_by_round["n"] = len(abstract_df.index)//200
abstract_df_by_round["Block"] = range(1,201)
sorted_abstract_df_by_round = abstract_df_by_round.iloc[:, [2, 3, 0, 1]]

expert_df_by_round = expert_df.groupby(["Round"]).agg({"Choice": "mean"})
expert_df_by_round = expert_df_by_round.groupby([np.arange(200)//5]).agg({"Choice": "mean"})
expert_df_by_round["Condition"] = "Description"
expert_df_by_round["n"] = len(expert_df.index)//200
expert_df_by_round["Block"] = range(1,201)
sorted_expert_df_by_round = expert_df_by_round.iloc[:, [2, 3, 0, 1]]
#
# attention_df.to_csv("Attention_check_passed.csv")
# sorted_expert_df_by_round.to_csv("Description_df.csv")
# sorted_abstract_df_by_round.to_csv("Experience_df.csv")

# print(sorted_expert_df_by_round.describe, sorted_abstract_df_by_round.describe)

    # present the plot
plt.plot(sorted_abstract_df_by_round["Block"], sorted_abstract_df_by_round["Choice"], color = "orange")
plt.plot(sorted_expert_df_by_round["Block"], sorted_expert_df_by_round["Choice"], color = "blue")

plt.ylim(0,1)
plt.title("Risk rates across 200 trials, Exp 1.")
plt.legend(["Experience (n=135)", "Description (n=109)"])
plt.xlabel("Block")
plt.ylabel("Risk rate")

    # plt.savefig("plot.png")
        #save the plot as a picture

plt.show()
# print(expert_df_by_round.loc[[0,1,99,100,101,102]])